package com.boge.ai.utils;

import com.boge.ai.entity.milvus.MilvusEntity;
import com.boge.ai.entity.rag.FaqItem;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.milvus.client.MilvusClient;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * 需要实现 RAG 的 向量数据库的工具类
 */
public class MilvusUtils {

    private MilvusServiceClient client;

    private static final ObjectMapper objectMapper = new ObjectMapper();

    public MilvusUtils() {
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withHost("localhost")
                .withPort(19530)
                .build();
        client = new MilvusServiceClient(connectParam);
    }

    /**
     * 创建数据集【类似于MYSQL数据库中的表结构】
     */
    public  void createCollection(){

        // 对应的数据集的字段集合
        List<FieldType> list = new ArrayList<>();
        list.add(FieldType
                .newBuilder()
                .withName("id")
                .withDescription("主键")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(true)
                .build());
        list.add(FieldType
                .newBuilder()
                .withName(MilvusEntity.Field.FEATURE)
                .withDescription("向量")
                .withDataType(DataType.FloatVector)
                .withDimension(MilvusEntity.FEATURE_DIM)
                .build());

        list.add(FieldType
                .newBuilder()
                .withName(MilvusEntity.Field.INSTRUCTION)
                .withDescription("向量对应的原文")
                .withDataType(DataType.VarChar)
                .withTypeParams(Collections.singletonMap("max_length","65535"))
                .build());

        list.add(FieldType
                .newBuilder()
                .withName(MilvusEntity.Field.OUTPUT)
                .withDescription("问题的答案")
                .withDataType(DataType.VarChar)
                .withTypeParams(Collections.singletonMap("max_length","65535"))
                .build());
        // 数据集的相关信息
        CreateCollectionParam createCollectionParam = CreateCollectionParam
                .newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withDescription("rag 案例")
                .withShardsNum(MilvusEntity.SHARDS_NUM)
                .withFieldTypes(list) // 设置对应的 字段信息
                .build();

        // 创建对应的 数据集
        client.createCollection(createCollectionParam);
        // 创建对应的索引
        createVectorIndex();
    }

    public  void createVectorIndex(){
        // 创建索引的参数
        CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withIndexName(MilvusEntity.Field.FEATURE)
                .withFieldName(MilvusEntity.Field.FEATURE)
                .withIndexType(IndexType.IVF_FLAT)
                .withMetricType(MetricType.L2) // 制定度量类型
                .withExtraParam("{\"nlist\":128}") // 根据索引类型提供额外参数，比如nlist
                .build();
        client.createIndex(createIndexParam);
    }

    /**
     * 加载数据到Redis服务中
     * @throws IOException
     */
    public void loadFaqDataToMilvus() throws IOException {
        InputStream inputStream = getClass().getClassLoader().getResourceAsStream("train_zh.json");
        if (inputStream == null) {
            throw new FileNotFoundException("File not found in resources: data.json");
        }

        try (BufferedReader reader = new BufferedReader(
                new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {

            String line;
            int index = 0;
            while ((line = reader.readLine()) != null) {
                if (line.trim().isEmpty()) continue;
                FaqItem item = objectMapper.readValue(line, FaqItem.class);
                // 获取到对应的instruction 做对应的Embedding
                String instruction = item.getInstruction();
                float[] embedding = LLMUtils.embedding(instruction, null);
                List<Float> floats = new ArrayList<>();
                for (float v : embedding) {
                    floats.add(v);
                }
                // 数组到集合的转换
                insertVectorData(floats,instruction,item.getOutput());
            }
            System.out.println("输出插入成功");
        }
    }

    /**
     * 向 向量数据库 添加数据
     */
    public  void insertVectorData(List<Float> vector,String instruction,String output){

        // id的数据
        List<List<Float>> floats = new ArrayList<>();
        floats.add(vector);

        // 添加需要添加的数据的集合
        List<InsertParam.Field> list = new ArrayList<>();
        list.add(new InsertParam.Field(MilvusEntity.Field.FEATURE,floats));
        list.add(new InsertParam.Field(MilvusEntity.Field.INSTRUCTION,Arrays.asList(instruction)));
        list.add(new InsertParam.Field(MilvusEntity.Field.OUTPUT,Arrays.asList(output)));

        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withFields(list)
                .build();
        client.insert(insertParam);
    }

    /**
     * 向量数据的查询
     */
    public  SearchResultsWrapper searchVectorData(List<Float> searchVector){
        // 定义的向量数据集
        List<List<Float>> searchVectorList = new ArrayList<>();
        searchVectorList.add(searchVector);
        // 查询的参数对象
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withMetricType(MetricType.L2)
                .withTopK(3) // 返回最接近的前三个记录
                .withVectorFieldName(MilvusEntity.Field.FEATURE)
                .withVectors(searchVectorList)
                .withOutFields(Arrays.asList(MilvusEntity.Field.ID,MilvusEntity.Field.OUTPUT,MilvusEntity.Field.INSTRUCTION)) // 定义需要返回的字段
                .build();

        SearchResults data = client.search(searchParam).getData();
        if(data != null){
            SearchResultsWrapper resultsWrapper = new SearchResultsWrapper(data.getResults());
            return resultsWrapper;
        }
        return null;
    }
}
